import pandas as pd
import os
import sys

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 导入项目模块
from data.fetcher import FinancialDataFetcher
from data.processor import FinancialDataProcessor

def fix_processor_issue():
    """修复数据处理器对非转置格式数据处理的问题"""
    print("\n===== 修复数据处理器对非转置格式数据处理的问题 =====")
    
    # 获取600900的原始数据
    fetcher = FinancialDataFetcher()
    stock_code = "600900"
    report_type = "年报"
    
    print(f"获取股票 {stock_code} 的{report_type}数据...")
    raw_data = fetcher.fetch_balance_sheet(stock_code, report_type)
    
    if raw_data is None or raw_data.empty:
        print("获取数据失败")
        return
    
    print(f"成功获取数据! 数据形状: {raw_data.shape}")
    
    # 修改processor.py中的categorize_data方法
    print("\n修改数据处理器的categorize_data方法...")
    
    # 创建修复后的处理器类
    class FixedFinancialDataProcessor(FinancialDataProcessor):
        def categorize_data(self, raw_data):
            """修复后的分类方法，支持非转置格式数据"""
            if raw_data is None or raw_data.empty:
                return None
            
            # 打印原始数据的基本信息，帮助调试
            print(f"原始数据形状: {raw_data.shape}")
            print(f"原始数据列名: {raw_data.columns[:5]}...")
            print(f"原始数据索引: {raw_data.index[:5]}...")
            
            # 检查数据结构，判断是否为转置格式
            is_transposed = True  # 默认假设是转置格式（列为日期，行为科目）
            
            # 检查列名是否为日期格式
            date_format_columns = False
            if len(raw_data.columns) > 0:
                # 检查前几列是否为日期格式
                for col in raw_data.columns[:3]:  # 检查前3列
                    col_str = str(col)
                    # 日期格式可能是20210331或2021-03-31
                    if col_str.replace('-', '').isdigit() and len(col_str.replace('-', '')) >= 8:
                        date_format_columns = True
                        break
                
                if not date_format_columns:
                    # 如果列名不是日期格式，可能不是转置格式
                    is_transposed = False
                    print("检测到非转置格式数据（列名不是日期）")
            
            # 检查第一行是否包含科目名称
            subject_in_first_row = False
            if raw_data.shape[0] > 0:
                first_row_values = [str(val).strip() for val in raw_data.iloc[0].values]
                common_subjects = ["流动资产", "货币资金", "应收账款", "存货", "固定资产", "短期借款"]
                if any(subj in val for val in first_row_values for subj in common_subjects):
                    subject_in_first_row = True
                    is_transposed = False
                    print("检测到非转置格式数据（第一行包含科目名称）")
            
            # 检查索引是否为科目名称
            subject_in_index = False
            if raw_data.index.dtype == 'object' and not raw_data.index.empty:
                # 检查索引的前几个值是否为科目名称
                sample_indices = [str(idx).strip() for idx in raw_data.index[:10].tolist()]
                common_subjects = ["流动资产", "货币资金", "应收账款", "存货", "固定资产", "短期借款"]
                if any(any(subj in idx for subj in common_subjects) for idx in sample_indices):
                    subject_in_index = True
                    is_transposed = True
                    print("确认转置格式数据（索引为科目名称）")
            
            # 综合判断数据格式
            if subject_in_index:
                is_transposed = True
            elif subject_in_first_row or not date_format_columns:
                is_transposed = False
            
            print(f"数据格式判断: {'转置格式' if is_transposed else '非转置格式'}")
            
            # 获取最新日期列或行
            latest_date = None
            if is_transposed:
                # 如果是转置格式，尝试找到最新的日期列
                date_columns = []
                for col in raw_data.columns:
                    col_str = str(col)
                    # 尝试识别日期格式（20210331或2021-03-31）
                    if col_str.replace('-', '').isdigit() and len(col_str.replace('-', '')) >= 8:
                        date_columns.append(col_str)
                
                if date_columns:
                    # 按日期排序，选择最新的日期
                    latest_date = sorted(date_columns)[-1]
                    print(f"找到最新日期列: {latest_date}")
                elif len(raw_data.columns) > 0:
                    # 如果没有找到日期列，使用第一列
                    latest_date = str(raw_data.columns[0])
                    print(f"未找到日期列，使用第一列: {latest_date}")
                else:
                    # 如果没有列，使用当前日期
                    from datetime import datetime
                    latest_date = datetime.now().strftime('%Y%m%d')
                    print(f"未找到有效列，使用当前日期: {latest_date}")
            else:
                # 如果是非转置格式，尝试从第一行获取日期
                if '报告日' in raw_data.columns:
                    latest_date = str(raw_data.iloc[0]['报告日'])
                    print(f"从'报告日'列获取日期: {latest_date}")
                elif '报告日期' in raw_data.columns:
                    latest_date = str(raw_data.iloc[0]['报告日期'])
                    print(f"从'报告日期'列获取日期: {latest_date}")
                else:
                    # 如果没有找到日期列，使用当前日期
                    from datetime import datetime
                    latest_date = datetime.now().strftime('%Y%m%d')
                    print(f"未找到日期列，使用当前日期: {latest_date}")
            
            # 添加报表日期列
            from datetime import datetime
            report_date = latest_date if latest_date else datetime.now().strftime('%Y%m%d')
            
            # 创建结果DataFrame，确保有一行数据
            result = pd.DataFrame(index=[0])
            result['报表日期'] = report_date
            
            # 按自定义分类整理数据
            for category, category_info in self.category_mapping.items():
                # 获取该分类下的科目列表
                if isinstance(category_info, dict):
                    items = category_info.get('items', [])
                else:
                    items = category_info  # 兼容旧版映射格式
                
                # 初始化该分类的总和为0.0
                result[category] = 0.0
                
                # 累加该分类下的所有科目
                for item in items:
                    # 检查科目是否在原始数据中
                    if is_transposed:
                        # 转置格式处理逻辑（索引为科目）
                        # 尝试精确匹配和部分匹配
                        matching_indices = [idx for idx in raw_data.index if str(idx).strip() == item]
                        if not matching_indices:
                            # 尝试部分匹配
                            matching_indices = [idx for idx in raw_data.index if item in str(idx).strip()]
                        
                        for idx in matching_indices:
                            try:
                                # 如果是转置格式，取最新日期的数据
                                if latest_date in raw_data.columns:
                                    amount = raw_data.loc[idx, latest_date]
                                    print(f"科目 {idx} 在 {latest_date} 的金额: {amount}")
                                elif len(raw_data.columns) > 0:
                                    # 如果最新日期不在列中，取第一列数据
                                    first_col = raw_data.columns[0]
                                    amount = raw_data.loc[idx, first_col]
                                    print(f"科目 {idx} 在 {first_col} 的金额: {amount}")
                                else:
                                    print(f"警告: 科目 {idx} 没有可用的数据列")
                                    continue
                                
                                # 如果金额是Series，取第一个值
                                if isinstance(amount, pd.Series):
                                    amount = amount.iloc[0]
                                
                                # 转换为浮点数并处理NaN
                                try:
                                    if pd.isna(amount) or amount == '':
                                        amount_float = 0.0
                                    else:
                                        # 处理可能的字符串格式（如带逗号的数字）
                                        if isinstance(amount, str):
                                            amount = amount.replace(',', '')
                                        amount_float = float(amount)
                                except (ValueError, TypeError) as e:
                                    print(f"警告: 无法将 {amount} 转换为浮点数，使用0代替: {e}")
                                    amount_float = 0.0
                                
                                # 累加到对应分类
                                result.at[0, category] += amount_float
                                print(f"累加 {idx} 到 {category}: +{amount_float}, 当前总和: {result.at[0, category]}")
                            except Exception as e:
                                print(f"处理科目 {idx} 时出错: {e}")
                    else:
                        # 非转置格式处理逻辑（列为科目）
                        if item in raw_data.columns:
                            try:
                                # 取第一行的数据（最新数据）
                                amount = raw_data.iloc[0][item]
                                print(f"科目 {item} 的金额: {amount}")
                                
                                # 转换为浮点数并处理NaN
                                try:
                                    if pd.isna(amount) or amount == '':
                                        amount_float = 0.0
                                    else:
                                        # 处理可能的字符串格式（如带逗号的数字）
                                        if isinstance(amount, str):
                                            amount = amount.replace(',', '')
                                        amount_float = float(amount)
                                except (ValueError, TypeError) as e:
                                    print(f"警告: 无法将 {amount} 转换为浮点数，使用0代替: {e}")
                                    amount_float = 0.0
                                
                                # 累加到对应分类
                                result.at[0, category] += amount_float
                                print(f"累加 {item} 到 {category}: +{amount_float}, 当前总和: {result.at[0, category]}")
                            except Exception as e:
                                print(f"处理科目 {item} 时出错: {e}")
            
            # 检查结果是否有非零值
            non_zero_cols = [col for col in result.columns if col != '报表日期' and result[col].sum() > 0]
            if non_zero_cols:
                print(f"成功分类! 包含非零值的分类数量: {len(non_zero_cols)} / {len(result.columns) - 1}")
            else:
                print("警告: 分类后的数据全为零，请检查数据处理逻辑")
            
            return result
    
    # 使用修复后的处理器处理数据
    processor = FixedFinancialDataProcessor()
    print("\n使用修复后的处理器处理数据...")
    categorized_data = processor.categorize_data(raw_data)
    
    if categorized_data is None or categorized_data.empty:
        print("数据处理失败")
        return
    
    print(f"\n分类后的数据:")
    print(categorized_data)
    
    # 保存分类后的数据
    output_path = "600900_fixed_categorized.csv"
    categorized_data.to_csv(output_path, encoding='utf-8-sig', index=False)
    print(f"\n分类后的数据已保存到: {output_path}")
    
    print("\n===== 修复建议 =====")
    print("1. 更新data/processor.py中的categorize_data方法，增加对非转置格式数据的处理")
    print("2. 确保科目映射关系中包含所有可能的科目名称")

if __name__ == "__main__":
    fix_processor_issue()